# 数据库配置和连接管理

from tortoise import Tortoise
from tortoise.contrib.fastapi import register_tortoise
from fastapi import FastAPI
from typing import Optional

from src.config.settings import get_settings
from src.constants import DATABASE_MODELS as DB_MODELS
from src.utils.logger import app_logger as logger


class DatabaseManager:
    """数据库管理器"""

    def __init__(self):
        self.db_url: Optional[str] = None
        self.is_connected: bool = False

    def get_database_url(self) -> str:
        """获取数据库连接URL"""
        config = get_settings()
        self.db_url = f"mysql://{config.MYSQL_USER}:{config.MYSQL_PASS}@{config.MYSQL_HOST}:{config.MYSQL_PORT}/{config.MYSQL_DB}"
        return self.db_url

    def get_tortoise_config(self) -> dict:
        """获取 Tortoise ORM 配置"""
        return {
            "connections": {"default": self.get_database_url()},
            "apps": {
                "models": {
                    "models": DB_MODELS + ["aerich.models"],
                    "default_connection": "default",
                }
            },
            "use_tz": False,
            "timezone": "Asia/Shanghai",
        }

    async def init_database(self):
        """初始化数据库连接"""
        try:
            config_dict = self.get_tortoise_config()
            await Tortoise.init(config=config_dict)
            self.is_connected = True

            # 获取连接信息用于日志
            db_url = config_dict["connections"]["default"]
            # 隐藏密码信息
            safe_url = db_url.split("@")[1] if "@" in db_url else db_url
            logger.info(f"数据库连接初始化成功: {safe_url}")

        except Exception as e:
            self.is_connected = False
            logger.error(f"数据库连接初始化失败: {e}")
            raise

    async def close_database(self):
        """关闭数据库连接"""
        if self.is_connected:
            await Tortoise.close_connections()
            self.is_connected = False
            logger.info("Database connections closed")

    async def generate_schemas(self, safe: bool = True):
        """生成数据库表结构"""
        try:
            await Tortoise.generate_schemas(safe=safe)
            logger.info("Database schemas generated successfully")
        except Exception as e:
            logger.error(f"Failed to generate schemas: {e}")
            raise


# 全局数据库管理器实例
db_manager = DatabaseManager()


def init_database_with_app(app: FastAPI):
    """为 FastAPI 应用初始化数据库"""
    register_tortoise(
        app,
        config=db_manager.get_tortoise_config(),
        generate_schemas=True,  # 自动生成表结构
        add_exception_handlers=True,  # 添加异常处理器
    )
    logger.info("Database registered with FastAPI app")
